import cv2
import glob
import os
from itertools import zip_longest
import re


def resize(img_array, align_mode):
    """调整图片大小，使所有图片的尺寸一致"""
    _height = len(img_array[0])
    _width = len(img_array[0][0])
    for i in range(1, len(img_array)):
        img = img_array[i]
        height = len(img)
        width = len(img[0])
        if align_mode == 'smallest':
            if height < _height:
                _height = height
            if width < _width:
                _width = width
        else:
            if height > _height:
                _height = height
            if width > _width:
                _width = width

    for i in range(0, len(img_array)):
        img1 = cv2.resize(img_array[i], (_width, _height), interpolation=cv2.INTER_CUBIC)
        img_array[i] = img1

    return img_array, (_width, _height)


def images_to_video(input_path, output_path):
    """将交错排列的图片合成视频"""
    img_array = []

    # 获取 num_t.jpg 和 num_ll.jpg 的文件路径
    def natural_sort_key(s):
        return [int(text) if text.isdigit() else text.lower() for text in re.split(r'(\d+)', s)]

    # 获取并排序图片
    t_images = sorted(glob.glob(os.path.join(input_path, '*_t.jpg')), key=natural_sort_key)
    ll_images = sorted(glob.glob(os.path.join(input_path, '*_ll.jpg')), key=natural_sort_key)

    # 检查是否有图片
    if not t_images and not ll_images:
        print(f"No valid images found in {input_path}")
        return

    # 交错排列图片
    for t_img, ll_img in zip_longest(t_images, ll_images):  # 使用 zip_longest 以支持长度不等的列表
        if t_img:
            t_image = cv2.imread(t_img)
            if t_image is None:
                print(f"{t_img} is error!")
            else:
                img_array.append(t_image)
        if ll_img:
            ll_image = cv2.imread(ll_img)
            if ll_image is None:
                print(f"{ll_img} is error!")
            else:
                img_array.append(ll_image)
    # 图片的大小需要一致
    img_array, size = resize(img_array, 'largest')
    fps = 25

    # 创建输出路径的文件夹（如果不存在）
    os.makedirs(output_path, exist_ok=True)

    # 输出视频文件路径
    video_path = os.path.join(output_path, 'fusion_demo.avi')
    out = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*'DIVX'), fps, size)

    for img in img_array:
        out.write(img)
    out.release()
    print(f"Video saved to {video_path}")


def main():
    # 指定输入路径和输出路径
    input_path = "/data1/lkf24/code/SpikeYOLO_old/SpikeYOLO/runs/detect/test_rescue_Forest2_yolov8m"  # 替换为你的图片路径
    output_path = input_path + '/video'  # 替换为你的视频输出路径

    # 检查输入路径是否存在
    if not os.path.exists(input_path):
        print(f"Input path {input_path} does not exist.")
        return

    # 合成视频
    images_to_video(input_path, output_path)


if __name__ == "__main__":
    main()